"""
数据库操作重试装饰器
用于处理SQLite等数据库的锁定问题
"""
import time
import functools
from sqlalchemy.exc import OperationalError
from api.utils.logger import get_logger

logger = get_logger("db_retry")


def retry_on_db_lock(max_retries=3, delay=0.5, backoff=2):
    """
    数据库操作重试装饰器
    
    Args:
        max_retries: 最大重试次数
        delay: 初始延迟时间（秒）
        backoff: 延迟时间的倍数因子
    
    Usage:
        @retry_on_db_lock(max_retries=3)
        def some_db_operation(db):
            db.query(...).update(...)
            db.commit()
    """
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            retry_delay = delay
            last_exception = None
            
            for attempt in range(max_retries + 1):
                try:
                    return func(*args, **kwargs)
                except OperationalError as e:
                    last_exception = e
                    error_msg = str(e).lower()
                    
                    # 检查是否是数据库锁定错误
                    if 'database is locked' in error_msg or 'locked' in error_msg:
                        if attempt < max_retries:
                            time.sleep(retry_delay)
                            retry_delay *= backoff
                            continue
                        else:
                            logger.error(
                                f"数据库锁定，已达最大重试次数: {func.__name__}"
                            )
                    # 不是锁定错误，直接抛出
                    raise
                except Exception as e:
                    # 其他异常直接抛出
                    raise
            
            # 如果所有重试都失败，抛出最后一个异常
            if last_exception:
                raise last_exception
        
        return wrapper
    return decorator


def retry_on_db_lock_async(max_retries=3, delay=0.5, backoff=2):
    """
    异步数据库操作重试装饰器
    
    Args:
        max_retries: 最大重试次数
        delay: 初始延迟时间（秒）
        backoff: 延迟时间的倍数因子
    """
    def decorator(func):
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            import asyncio
            retry_delay = delay
            last_exception = None
            
            for attempt in range(max_retries + 1):
                try:
                    return await func(*args, **kwargs)
                except OperationalError as e:
                    last_exception = e
                    error_msg = str(e).lower()
                    
                    # 检查是否是数据库锁定错误
                    if 'database is locked' in error_msg or 'locked' in error_msg:
                        if attempt < max_retries:
                            await asyncio.sleep(retry_delay)
                            retry_delay *= backoff
                            continue
                        else:
                            logger.error(
                                f"数据库锁定，已达最大重试次数: {func.__name__}"
                            )
                    # 不是锁定错误，直接抛出
                    raise
                except Exception as e:
                    # 其他异常直接抛出
                    raise
            
            # 如果所有重试都失败，抛出最后一个异常
            if last_exception:
                raise last_exception
        
        return wrapper
    return decorator